#include "StdAfx.h"
#include "AESMath.h"

/***
 *  Public Interface
 */

AESMath::AESMath(void)
{
	newMathTables();
}

AESMath::~AESMath(void)
{
	deleteMathTables();
}


// Routines to access table entries
byte AESMath::SBox(byte b) {
	
	return _tblSubBytes[b & 0xff];
}

byte AESMath::invSBox(byte b) {
	return _tblSubBytesInverse[b & 0xff];
}


byte AESMath::Rcon(int i) {
	return _tblPowX[i-1];
}


/***
 *     Iniitalizers and the Loaders for the Tables
 *
 */

void AESMath::newMathTables(void)
{
	/***
	 *  TODO: will implement exception handling, so we can 
	 *  write readable code that handles multiple allocations
	 *  and error possibilities.
	 */
	_tblExponentials = new byte[AES_GFMATH_TABLE_LENGTH];
	_tblLogarithms = new byte[AES_GFMATH_TABLE_LENGTH];; 
	_tblSubBytes = new byte[AES_GFMATH_TABLE_LENGTH];;		  
	_tblSubBytesInverse = new byte[AES_GFMATH_TABLE_LENGTH];;  
	_tblMultInverse = new byte[AES_GFMATH_TABLE_LENGTH];;	  
	_tblPowX = new byte[AES_GFMATH_TABLE_LENGTH];			

	if ( !_tblExponentials || !_tblLogarithms || 
		 !_tblSubBytes || !_tblSubBytesInverse  || 
		 !_tblMultInverse || !_tblPowX)
	{
		perror("newMathTables Fatal error: no memory available");
	}

}

void AESMath::deleteMathTables(void)
{
	delete [] _tblExponentials;
	delete [] _tblLogarithms; 
	delete [] _tblSubBytes;		  
	delete [] _tblSubBytesInverse;  
	delete [] _tblMultInverse;	  
	delete [] _tblPowX;			

	_tblExponentials = _tblLogarithms = _tblSubBytes = _tblSubBytesInverse = _tblMultInverse = _tblPowX = NULL;
}

// create and load the Exponentials table
void AESMath::loadExponentials() 
{
	byte x			= (byte) AES_FIELD_VALUE_01;
	int index		= 0;
	int tblLength	= AES_GFMATH_TABLE_LENGTH - 1;

	// returns 1
	_tblExponentials[index++] = (byte) AES_FIELD_VALUE_01;
	
	for (int i = 0; i < tblLength; i++) 
	{
		// raising by successive power of 3
		byte y = FFMul(x, (byte) AES_EXPONENTIAL_BASE_03);

		// assign to table
		_tblExponentials[index++] = y;

		// get ready for the next
		x = y;
	}
}


// load the Logarigthms table using the Exponentials table
void AESMath::loadLogarithms() 
{ 
	int tblLength	= AES_GFMATH_TABLE_LENGTH - 1;

	// note this uses upto 254
	for (int i = 0; i < tblLength; i++) 
	{
		_tblLogarithms[_tblExponentials[i] & 0xff] = (byte)i;
	}
}


// load in the table for multiplicative inverse
void AESMath::loadInv() 
{
	int tblLength	= AES_GFMATH_TABLE_LENGTH;
	
	for (int i = 0; i < tblLength; i++)
	{
		_tblMultInverse[i] = (byte)(FFInv((byte)(i & 0xff)) & 0xff);
	}
}


// load in the table Sbox
void AESMath::loadSbox() 
{
	int tblLength	= AES_GFMATH_TABLE_LENGTH;

	for (int i = 0; i < tblLength; i++)
	{
		_tblSubBytes[i] = (byte)(subBytes((byte)(i & 0xff)) & 0xff);
	}
}

// load the InverseSBox table using the SBox table
void AESMath::loadInverseSBox() 
{
	int tblLength	= AES_GFMATH_TABLE_LENGTH;

	for (int i = 0; i < tblLength; i++) 
	{
		_tblSubBytesInverse[_tblSubBytes[i] & 0xff] = (byte)i;
	}
}


// loadPowX: load the _tblPowX table using multiplication
void AESMath::loadPowX() 
{
	byte x = (byte) AES_EXPONENTIAL_BASE_02;

	byte xp = x;

	_tblPowX[0] = 1; _tblPowX[1] = x;

	for (int i = 2; i < 15; i++) 
	{
		xp = FFMul(xp, x);
		_tblPowX[i] = xp;
	}
}

	// FFInv: the multiplicative inverse of a byte value
byte AESMath::FFInv(byte b) 
{
	byte e = _tblLogarithms[b & 0xff];

	return _tblExponentials[0xff - (e & 0xff)];
}

//  return the ith bit of a byte
int AESMath::ithBit(byte b, int i) 
{
	int m[] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80};

	return (b & m[i]) >> i;
}

// SubBytes Transformation 
//   - a reversible transformation
//   - Transforms each byte in the plaintext into another byte.
int AESMath::subBytes(byte b) 
{
	byte	inB = b;
	int		res = 0;

	if (b != 0) { // if b == 0, leave it alone
		b = (byte)(FFInv(b) & 0xff);
	}
	
	byte c = (byte) 0x63;
	for (int i = 0; i < 8; i++) 
	{
		int temp = 0;

		temp = ithBit(b, i) ^ ithBit(b, (i+4)%8) ^ ithBit(b, (i+5)%8) ^ 
			ithBit(b, (i+6)%8) ^ ithBit(b, (i+7)%8) ^ ithBit(c, i);

		res = res | (temp << i);
	}
	
	return res;
}
// FFMulFast: fast multiply using table lookup
byte AESMath::FFMulFast(byte a, byte b)
{
	int t = 0;;

	if (a == 0 || b == 0) {
		return 0;
	}

	t = (_tblLogarithms[(a & 0xff)] & 0xff) + (_tblLogarithms[(b & 0xff)] & 0xff);

	// if value is > 255, this reduction finds the correct value in the finite set
	if (t > 255) 
	{
		t = t - 255;
	}
	return _tblExponentials[(t & 0xff)];
}

// FFMul: slow multiply, using shifting
byte AESMath::FFMul(byte a, byte b) {
	byte aa = a, bb = b, r = 0, t;
	
	while (aa != 0) 
	{
		if ((aa & 1) != 0)
		{
			r = (byte)(r ^ bb);
		}

		t = (byte)(bb & 0x80);

		bb = (byte)(bb << 1);
		if (t != 0){
			bb = (byte)(bb ^ 0x1b);
		}

		aa = (byte)((aa & 0xff) >> 1);
	}
	
	return r;
}
